"""
python p2_solution_deap_genetic_4t_gpt.py
"""

import numpy as np
import random
from deap import base, creator, tools, algorithms
from datetime import datetime
import json
import read_4_json as r4j

# Load data from your provided module
data_warehouse = r4j.load_warehouse_data()
data_category_inventory = r4j.load_average_inventory_data()
data_category_sales = r4j.load_average_sales_data()
data_category_association = r4j.load_association_data()
mx_category_association = r4j.create_adjacency_matrix()

NUM_W = r4j.NUM_W
NUM_C = r4j.NUM_C

# Fitness evaluation function
def evaluate(individual):
    f1, f2, f3, f4, fitness_value = cal_4f_fitness(np.array(individual))
    return fitness_value,

# Define individual and population
creator.create("FitnessMax", base.Fitness, weights=(1.0,))  # Maximize the fitness
creator.create("Individual", list, fitness=creator.FitnessMax)

toolbox = base.Toolbox()
toolbox.register("attr_warehouse", random.randint, 0, NUM_W - 1)
toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_warehouse, NUM_C)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

# Genetic algorithm functions
toolbox.register("evaluate", evaluate)
toolbox.register("mate", tools.cxUniform, indpb=0.05)
toolbox.register("mutate", tools.mutUniformInt, low=0, up=NUM_W - 1, indpb=0.1)
toolbox.register("select", tools.selTournament, tournsize=3)

# Global variable to track best solution
absolute_best_individual = None
absolute_best_fitness = float("-inf")

# Logging best solution
def log_best_solution(individual):
    global absolute_best_individual, absolute_best_fitness
    fitness_value = toolbox.evaluate(individual)[0]
    
    if fitness_value > absolute_best_fitness:
        absolute_best_individual = individual[:]
        absolute_best_fitness = fitness_value
        
        # Save best solution
        solution = [{"category_id": i, "warehouse_id": wid} for i, wid in enumerate(individual)]
        with open("../fujian/p2/deap_genetic_4t/best_ever/best_ever_solution.json", "w") as f:
            json.dump(solution, f, indent=4)
        
        # Log best fitness and objectives
        f1, f2, f3, f4, _ = cal_4f_fitness(np.array(individual))
        with open("../fujian/p2/deap_genetic_4t/best_ever/best_ever_individual.log", "a") as log_file:
            log_file.write(f"{datetime.now()} - Best fitness: {fitness_value}, f1: {f1}, f2: {f2}, f3: {f3}, f4: {f4}\n")
        print(f"New best individual saved with fitness: {fitness_value}")

# Main genetic algorithm process
def main():
    random.seed(64)
    pop = toolbox.population(n=300)  # Population size

    NGEN = 5000  # Max number of generations
    CXPB, MUTPB = 0.5, 0.2  # Crossover and mutation probabilities

    for gen in range(NGEN):
        offspring = toolbox.select(pop, len(pop))
        offspring = list(map(toolbox.clone, offspring))

        # Apply crossover and mutation
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < CXPB:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values

        for mutant in offspring:
            if random.random() < MUTPB:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        # Evaluate fitness for invalid individuals
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = map(toolbox.evaluate, invalid_ind)
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit

        # Replace population with offspring
        pop[:] = offspring

        # Track best individual in each generation
        best_ind = tools.selBest(pop, 1)[0]
        log_best_solution(best_ind)

        # Elitism: inject absolute best into population to maintain progress
        if absolute_best_individual:
            pop[random.randint(0, len(pop) - 1)] = creator.Individual(absolute_best_individual)

if __name__ == "__main__":
    main()
